#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Simple WebSocket Client for Testing ASR WebSocket Server
简单的ASR WebSocket服务器测试客户端
"""

import asyncio
import websockets
import json
import base64
import sys
import os

# 添加src目录到路径中以便导入
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

import numpy as np


def create_test_audio_data():
    """
    创建测试音频数据
    :return: 音频数据字节串
    """
    # 创建1秒的静音音频数据 (16kHz, 16-bit)
    sample_rate = 16000
    duration = 1  # 1秒
    num_samples = sample_rate * duration
    
    # 生成静音音频数据
    audio_data = np.zeros(num_samples, dtype=np.int16)
    
    return audio_data.tobytes()


async def test_websocket_client():
    """
    测试WebSocket客户端
    """
    uri = "ws://localhost:8767"
    
    try:
        async with websockets.connect(uri) as websocket:
            print("已连接到WebSocket服务器")
            
            # 接收欢迎消息
            welcome_msg = await websocket.recv()
            print(f"收到服务器消息: {welcome_msg}")
            
            # 发送重置上下文消息
            reset_msg = {
                "type": "reset_context"
            }
            
            await websocket.send(json.dumps(reset_msg))
            print("已发送重置上下文消息")
            
            # 接收响应
            response = await websocket.recv()
            print(f"收到服务器响应: {response}")
            
            # 创建测试音频数据
            test_audio_data = create_test_audio_data()
            encoded_audio = base64.b64encode(test_audio_data).decode('utf-8')
            
            # 发送测试音频数据
            message = {
                "type": "audio",
                "data": encoded_audio
            }
            
            await websocket.send(json.dumps(message))
            print("已发送测试音频数据")
            
            # 接收处理结果
            response = await websocket.recv()
            print(f"收到识别结果: {response}")
            
            # 尝试发送第二个音频数据
            await websocket.send(json.dumps(message))
            print("已发送第二个测试音频数据")
            
            # 接收处理结果
            response = await websocket.recv()
            print(f"收到第二个识别结果: {response}")
            
    except websockets.exceptions.ConnectionRefusedError:
        print("连接被拒绝，请确保WebSocket服务器正在运行")
        print("启动服务器命令: python main.py --websocket --ws-port 8767")
    except Exception as e:
        print(f"测试客户端出错: {str(e)}")


def main():
    """
    主函数
    """
    print("ASR WebSocket服务简单测试")
    print("=" * 40)
    
    # 运行测试客户端
    asyncio.run(test_websocket_client())


if __name__ == "__main__":
    main()